package com.itheima.stargraph.service.impl;

import cn.hutool.core.bean.BeanUtil;
import cn.hutool.core.util.StrUtil;
import cn.hutool.json.JSONUtil;
import com.itheima.stargraph.dto.request.Text2ImageReqDto;
import com.itheima.stargraph.dto.respone.Text2ImageResDto;
import com.itheima.stargraph.pojo.ComfyuiModel;
import com.itheima.stargraph.pojo.ComfyuiRequestDto;
import com.itheima.stargraph.pojo.ComfyuiTask;
import com.itheima.stargraph.service.FreeMarkerService;
import com.itheima.stargraph.service.OllamaService;
import com.itheima.stargraph.service.Text2ImageService;
import com.itheima.stargraph.utils.UserThreadLocalUtils;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.data.redis.core.ZSetOperations;
import org.springframework.stereotype.Service;

import java.util.concurrent.TimeUnit;

/**
 * 文生图业务实现类
 *
 * @Author mr.wu
 * @Date 2025-6-11 11:44
 */
@Service
public class Text2ImageServiceImpl implements Text2ImageService {

    @Autowired
    private OllamaService ollamaService;

    @Autowired
    private FreeMarkerService freemarkerService;

    //comfyui服务端指定的客户端ID值（必须固定）
    private static final String CLIENT_ID = "star-graph";

    /**
     * 获取Comfyui任务实例
     *
     * @param dto
     * @return
     */
    @Override
    public ComfyuiTask getComfyuiTask(Text2ImageReqDto dto) {

        //1. 将dto复制freemarker的模型实例（ComfyuiModel实例）
        ComfyuiModel comfyuiModel = BeanUtil.toBean(dto, ComfyuiModel.class);

        //2. 将正向词和负向词翻译成英文（调用ollama接口）
        String prompt = ollamaService.translate(dto.getPrompt());//正向词翻译结果
        String reverse = ollamaService.translate(dto.getReverse());//负向词翻译结果

        //3. 补充设置更多属性的值
        comfyuiModel.setPrompt(prompt); //正向词
        comfyuiModel.setReverse(reverse); //负向词
        comfyuiModel.setModelName(dto.modelName()); //模型名
        comfyuiModel.setSamplerName(dto.samplerName());//采样器名称
        comfyuiModel.setScheduler(dto.scheduler());//调度器名
        comfyuiModel.setWidth(dto.width());//宽
        comfyuiModel.setHeight(dto.height());//高

        //4. 使用freemarker基于模型数据渲染工作流JSON
        String workflowJSON = freemarkerService.render(comfyuiModel);

        //5. 封装comfyui添加流程任务的参数实例（ComfyuiRequestDto实例）
        ComfyuiRequestDto comfyuiRequestDto = new ComfyuiRequestDto(CLIENT_ID, JSONUtil.parseObj(workflowJSON));

        //6. 封装需要添加到缓存队列中的任务实例（ComfyuiTask任务实例）
        ComfyuiTask comfyuiTask = new ComfyuiTask(UserThreadLocalUtils.getUser().getId(), dto.getClientId(),comfyuiRequestDto);

        return comfyuiTask;
    }

    /**
     * 任务的分布式ID计数器的redisKey （String类型）
     */
    private static final String REDIS_KEY_TASK_COUNTER = "star:graph:task:counter";

    /**
     * 任务队列的redisKey （Zset类型）
     */
    private static final String REDIS_KEY_TASK_QUEUE = "star:graph:task:queue";

    /**
     * 任务的分布式ID计数器的redisKey  （String类型）
     */
    private static final String REDIS_KEY_TASK_INFO = "star:graph:task:%s";

    @Autowired
    private StringRedisTemplate redisTemplate;


    /**
     * 将任务添加到缓存中排队
     *
     * @param comfyuiTask
     */
    @Override
    public void addComfyuiTask(ComfyuiTask comfyuiTask) {
        //1. 获取当前任务的执行序号（基于Redis分布式计数器自增1）
        Long index = redisTemplate.opsForValue().increment(REDIS_KEY_TASK_COUNTER);

        //2. 将任务存入Redis的ZSET队列中排队（value：任务id，score：执行序号）
        redisTemplate.opsForZSet().add(REDIS_KEY_TASK_QUEUE,comfyuiTask.getId(),index);

        //3. 将任务详情JSON存入Redis的String中
        redisTemplate.opsForValue().set(String.format(REDIS_KEY_TASK_INFO,comfyuiTask.getId()),JSONUtil.toJsonStr(comfyuiTask));

        //4. 补充设置任务的执行序号
        comfyuiTask.setIndex(index);
    }

    /**
     * 文生图
     *
     * @param dto
     * @return
     */
    @Override
    public Text2ImageResDto prompt(Text2ImageReqDto dto) {
        //1. 获取任务
        ComfyuiTask comfyuiTask = this.getComfyuiTask(dto);

        //2. 任务排队
        this.addComfyuiTask(comfyuiTask);

        //3. 封装响应结果
        Text2ImageResDto resDto = new Text2ImageResDto();
        resDto.setPid(comfyuiTask.getId());
        resDto.setQueueIndex(comfyuiTask.getIndex());

        return resDto;
    }

    /**
     * 获取并删除缓存队列中的任务
     *
     * @return
     */
    @Override
    public ComfyuiTask popComfyuiTask() {
        //1. 判断Redis的ZSET队列是否元素
        Long size = redisTemplate.opsForZSet().size(REDIS_KEY_TASK_QUEUE);
        if(size==0) {
            return null;
        }

        //2. 从Redis的ZSET队列中获取分值最小（执行优先级最高）的任务ID
        ZSetOperations.TypedTuple<String> tuple = redisTemplate.opsForZSet().popMin(REDIS_KEY_TASK_QUEUE);
        if(tuple==null) {
            return null;
        }
        String taskId = tuple.getValue();

        //3. 根据任务ID从Redis的String缓存中获取任务JSON
        String comfyuiTaskJSON = redisTemplate.opsForValue().get(String.format(REDIS_KEY_TASK_INFO, taskId));
        if(StrUtil.isBlank(comfyuiTaskJSON)) {
            return null;
        }

        //4. 将任务JSON解析为任务的对象实例
        return JSONUtil.toBean(comfyuiTaskJSON,ComfyuiTask.class);
    }


    /**
     * 为websocket协议推送消息给前端准备的缓存key，目的是存储任务详情
     */
    private static final String REDIS_KEY_WS_TASK = "star:graph:task:ws:%s";

    /**
     * 添加开始执行的任务
     *
     * @param promptId  任务ID
     * @param comfyuiTask 任务对象
     */
    @Override
    public void addStartedTask(String promptId, ComfyuiTask comfyuiTask) {
        //key就是promptId  value就是任务详情的JSON字符串
        redisTemplate.opsForValue().set(String.format(REDIS_KEY_WS_TASK, promptId),JSONUtil.toJsonStr(comfyuiTask), 30, TimeUnit.MINUTES);
    }

    /**
     * 根据任务ID获取任务详情
     *
     * @param promptId
     * @return
     */
    @Override
    public ComfyuiTask getStartedTask(String promptId) {
        String comfyuiTaskJSON = redisTemplate.opsForValue().get(String.format(REDIS_KEY_WS_TASK, promptId));
        if(StrUtil.isBlank(comfyuiTaskJSON)) {
            return null;
        }
        return JSONUtil.toBean(comfyuiTaskJSON,ComfyuiTask.class);
    }
}
